import os.path
import configparser
from project_path import get_project_path
import json


# 实例化配置文件
config = configparser.ConfigParser()


def load_llm() -> tuple:
    """
    获取 LLM 相关配置
    :return: URL，密钥，模型名称
    """
    with open(os.path.join(get_project_path(), 'config.ini'), 'r', encoding='utf-8') as config_file:
        config.read_file(config_file)
    base_url = config['LLM']['base_url']
    api_key = config['LLM']['api_key']
    model_id = config['LLM']['model_id']
    return base_url, api_key, model_id


def load_prompt(user_input=None, task=None, case_qa=None, sql=None, error_info=None, qa_level=None,
                sql_result=None) -> str:
    """
    组织系统提示词 system_prompt
    :param user_input: 用户输入
    :param task: 任务类型
    :param case_qa: (针对 text-to-sql 任务) 提示案例
    :param sql: (针对 rewrite sql 任务) 报错 SQL
    :param error_info: (针对 rewrite sql 任务) 执行 SQL 报错详情
    :param qa_level: (针对 common question 任务) 问题层级
    :param sql_result: (针对 sql to markdown 任务) SQL 执行结果
    :return: 系统提示词
    """
    # 获取配置文件
    with open(os.path.join(get_project_path(), 'config.ini'), 'r', encoding='utf-8') as config_file:
        config.read_file(config_file)
    # 获取数据库信息
    dialect = config['DATABASE']['dialect']
    with open(os.path.join(get_project_path(), 'data', 'db_schema.json'), 'r', encoding='utf-8') as f:
        db_schema = str(json.load(f))
    # 获取系统提示词
    if task == 'text2sql':
        # text-to-sql 任务提示词
        return config['PROMPT']['text2sql'].format(dialect=dialect, db_schema=db_schema, qa_pair=case_qa,
                                                   question=user_input)
    elif task == 'rewriteSql':
        # rewrite sql 任务提示词
        return config['PROMPT']['rewrite_sql'].format(dialect=dialect, question=user_input, sql=sql,
                                                      error_info=error_info, db_schema=db_schema, qa_pair=case_qa)
    elif task == 'commonQuestionMatch':
        # common question 任务提示词
        return config['PROMPT']['common_question_match'].format(qa_level=qa_level, question=user_input)
    elif task == 'statisQuestionMatch':
        # statis question 任务提示词
        return config['PROMPT']['statis_question_match'].format(question=user_input)
    elif task == 'sql2md':
        # sql to markdown 任务提示词
        return config['PROMPT']['sql2md'].format(question=user_input, sql=sql, sql_result=sql_result)
    elif task == 'expansion_question':
        # expansion question 任务提示词
        return config['PROMPT']['expansion_question'].format(db_schema=db_schema, qa_pair=case_qa, question=user_input)
    else:
        return config['PROMPT']['system_prompt']


def load_host_port() -> tuple:
    """
    获取 host 和 port 相关配置
    :return: host，port
    """
    with open(os.path.join(get_project_path(), 'config.ini'), 'r', encoding='utf-8') as config_file:
        config.read_file(config_file)
    host = config['APP']['host']
    port = int(config['APP']['port'])
    return host, port


def load_database() -> tuple:
    """
    获取数据库的 dialect 和 database
    :return: dialect，database
    """
    with open(os.path.join(get_project_path(), 'config.ini'), 'r', encoding='utf-8') as config_file:
        config.read_file(config_file)
    dialect = config['DATABASE']['dialect']
    database = config['DATABASE']['database']
    return dialect, database


if __name__ == '__main__':
    print(load_llm())
    print(load_prompt())
